/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 */

package async;

import org.apache.jmeter.control.Controller;
import org.apache.jmeter.control.GenericController;
import org.apache.jmeter.control.NextIsNullException;
import org.apache.jmeter.samplers.SampleEvent;
import org.apache.jmeter.samplers.SampleListener;
import org.apache.jmeter.samplers.SampleResult;
import org.apache.jmeter.samplers.Sampler;
import org.apache.jmeter.testelement.property.BooleanProperty;
import org.apache.jmeter.threads.*;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.io.Serializable;
import java.util.ArrayList;
import java.util.TreeMap;
import java.util.concurrent.CountDownLatch;

/**
 * Transaction Controller to measure transaction times
 * <p>
 * There are two different modes for the controller:
 * - generate additional total sample after nested samples (as in MyJMeter 2.2)
 * - generate parent sampler containing the nested samples
 */
public class MyTransactionController extends GenericController implements SampleListener, Controller, Serializable {

	/**
	 * Used to identify Transaction Controller Parent Sampler
	 */
	static final String NUMBER_OF_SAMPLES_IN_TRANSACTION_PREFIX = "Number of samples in transaction : ";
	private static final long serialVersionUID = 222234L;
	private static final String TRUE = Boolean.toString(true); // i.e. "true"
	private static final String GENERATE_PARENT_SAMPLE = "TransactionController.parent";// $NON-NLS-1$
	private static final String INCLUDE_TIMERS = "TransactionController.includeTimers";// $NON-NLS-1$
	private static final Logger log = LoggerFactory.getLogger(MyTransactionController.class);
	private static final boolean DEFAULT_VALUE_FOR_INCLUDE_TIMERS = true; // default true for compatibility
	public static TreeMap<String, MyTransactionController> myTransactionControllerMap = new TreeMap<String, MyTransactionController>();
	public CountDownLatch countDownLatch;
	/**
	 * Only used in parent Mode
	 */
	private transient MyTransactionSampler myTransactionSampler;

	/**
	 * Only used in NON parent Mode
	 */
	private transient ListenerNotifier lnf;

	/**
	 * Only used in NON parent Mode
	 */
	private transient SampleResult res;

	/**
	 * Only used in NON parent Mode
	 */
	private transient int calls;

	/**
	 * Only used in NON parent Mode
	 */
	private transient int noFailingSamples;

	/**
	 * Cumulated pause time to excluse timer and post/pre processor times
	 * Only used in NON parent Mode
	 */
	private transient long pauseTime;

	/**
	 * Previous end time
	 * Only used in NON parent Mode
	 */
	private transient long prevEndTime;
	private ArrayList<SampleResult> sampleResultList = new ArrayList<SampleResult>();

	/**
	 * Creates a Transaction Controller
	 */
	public MyTransactionController() {
		lnf = new ListenerNotifier();
	}

	/**
	 * @param res {@link SampleResult}
	 * @return true if res is the ParentSampler transactions
	 */
	public static boolean isFromTransactionController(SampleResult res) {
		return res.getResponseMessage() != null &&
				res.getResponseMessage().startsWith(
						MyTransactionController.NUMBER_OF_SAMPLES_IN_TRANSACTION_PREFIX);
	}

	@Override
	protected Object readResolve() {
		super.readResolve();
		lnf = new ListenerNotifier();
		return this;
	}

	/**
	 * @return {@code true} if a parent sample will be generated
	 */
	public boolean isGenerateParentSample() {
		return getPropertyAsBoolean(GENERATE_PARENT_SAMPLE);
	}

///////////////// Transaction Controller - parent ////////////////

	/**
	 * @param generateParent flag whether a parent sample should be generated.
	 */
	public void setGenerateParentSample(boolean generateParent) {
		setProperty(new BooleanProperty(GENERATE_PARENT_SAMPLE, generateParent));
	}

	/**
	 * @see Controller#next()
	 */
	@Override
	public Sampler next() {
		if (isGenerateParentSample()) {
			return nextWithTransactionSampler();
		}
		return nextWithoutTransactionSampler();
	}

////////////////////// Transaction Controller - additional sample //////////////////////////////

	private Sampler nextWithTransactionSampler() {
		// Check if transaction is done
		if (myTransactionSampler != null && myTransactionSampler.isTransactionDone()) {
			if (log.isDebugEnabled()) {
				log.debug("End of transaction {}", getName());
			}
			// This transaction is done
			myTransactionSampler = null;
			return null;
		}

		// Check if it is the start of a new transaction
		if (isFirst()) // must be the start of the subtree
		{
			if (log.isDebugEnabled()) {
				log.debug("Start of transaction {}", getName());
			}
			myTransactionSampler = new MyTransactionSampler(this, getName());
		}

		// Sample the children of the transaction
		Sampler subSampler = super.next();
		myTransactionSampler.setSubSampler(subSampler);
		// If we do not get any sub samplers, the transaction is done
		if (subSampler == null) {
			myTransactionSampler.setTransactionDone();
		}
		return myTransactionSampler;
	}

	@Override
	protected Sampler nextIsAController(Controller controller) throws NextIsNullException {
		if (!isGenerateParentSample()) {
			return super.nextIsAController(controller);
		}
		Sampler returnValue;
		Sampler sampler = controller.next();
		if (sampler == null) {
			currentReturnedNull(controller);
			// We need to call the super.next, instead of this.next, which is done in GenericController,
			// because if we call this.next(), it will return the myTransactionSampler, and we do not want that.
			// We need to get the next real sampler or controller
			returnValue = super.next();
		} else {
			returnValue = sampler;
		}
		return returnValue;
	}

	private Sampler nextWithoutTransactionSampler() {
		if (isFirst()) // must be the start of the subtree
		{
			log.info(Thread.currentThread().getName() + "..." + this.hashCode());
			MyTransactionController.myTransactionControllerMap.put(Thread.currentThread().getName(), this);
			countDownLatch = new CountDownLatch(super.subControllersAndSamplers.size());

			calls = 0;
			noFailingSamples = 0;
			res = new SampleResult();
			sampleResultList.add(res); // 这个也给加一下
			res.setSampleLabel(getName());
			// Assume success
			res.setSuccessful(true);
			res.sampleStart();
			prevEndTime = res.getStartTime();//???
			pauseTime = 0;
		}
		boolean isLast = current == super.subControllersAndSamplers.size();
		Sampler returnValue = super.next();
		if (returnValue == null && isLast) // Must be the end of the controller
		{

			try {
				countDownLatch.await();

				// 采样结束时间前置，更加准确
				res.sampleEnd();
			} catch (InterruptedException e) {
				// TODO Auto-generated catch block
				e.printStackTrace();
			}

			// 将采样发生后的动作后置到此处
			for (SampleResult sampleResult : sampleResultList) {
				res.setThreadName(sampleResult.getThreadName());
				res.setBytes(res.getBytesAsLong() + sampleResult.getBytesAsLong());
				res.setSentBytes(res.getSentBytes() + sampleResult.getSentBytes());
				if (!isIncludeTimers()) {// Accumulate waiting time for later
					pauseTime += sampleResult.getEndTime() - sampleResult.getTime() - prevEndTime;
					prevEndTime = sampleResult.getEndTime();
				}
				if (!sampleResult.isSuccessful()) {
					res.setSuccessful(false);
					noFailingSamples++;
				}
				res.setAllThreads(sampleResult.getAllThreads());
				res.setGroupThreads(sampleResult.getGroupThreads());
				res.setLatency(res.getLatency() + sampleResult.getLatency());
				res.setConnectTime(res.getConnectTime() + sampleResult.getConnectTime());
			}

			if (res != null) {
				// See BUG 55816
				if (!isIncludeTimers()) {
					long processingTimeOfLastChild = res.currentTimeInMillis() - prevEndTime;
					pauseTime += processingTimeOfLastChild;
				}
				// 使用异步之后这里计算有问题，直接改成0，不知道这个是什么东西
				// res.setIdleTime(pauseTime+res.getIdleTime());
				res.setIdleTime(0);
				res.setResponseMessage(MyTransactionController.NUMBER_OF_SAMPLES_IN_TRANSACTION_PREFIX + calls + " - number of failing samples : " + noFailingSamples);
				if (res.isSuccessful()) {
					res.setResponseCodeOK();
				}
				notifyListeners();
			}

			// 此处可能存在性能问题，使用了同步代码块
			for (SampleResult result : sampleResultList) {
				MyJtlLogWriter.print(result);
			}

			// 听说只要没有抛出异常，try-catch对性能没什么损耗
			try {
				MyJtlLogWriter.writer.flush();
			} catch (IOException e) {
				log.error(e.getMessage(), e);
			}

			// 清理以供下一次使用
			sampleResultList.clear();
		} else {
			// We have sampled one of our children
			calls++;
		}

		return returnValue;
	}

	/**
	 * @see GenericController#triggerEndOfLoop()
	 */
	@Override
	public void triggerEndOfLoop() {
		if (!isGenerateParentSample()) {
			if (res != null) {
				res.setIdleTime(pauseTime + res.getIdleTime());
				res.sampleEnd();
				res.setSuccessful(TRUE.equals(JMeterContextService.getContext().getVariables().get(JMeterThread.LAST_SAMPLE_OK)));
				res.setResponseMessage(MyTransactionController.NUMBER_OF_SAMPLES_IN_TRANSACTION_PREFIX + calls + ", number of failing samples : " + noFailingSamples);
				notifyListeners();
			}
		} else {
			Sampler subSampler = myTransactionSampler.getSubSampler();
			// See Bug 56811
			// triggerEndOfLoop is called when error occurs to end Main Loop
			// in this case normal workflow doesn't happen, so we need
			// to notify the childs of TransactionController and
			// update them with SubSamplerResult
			if (subSampler instanceof MyTransactionSampler) {
				MyTransactionSampler tc = (MyTransactionSampler) subSampler;
				myTransactionSampler.addSubSamplerResult(tc.getTransactionResult());
			}
			myTransactionSampler.setTransactionDone();
			// This transaction is done
			myTransactionSampler = null;
		}
		super.triggerEndOfLoop();
	}

	/**
	 * Create additional SampleEvent in NON Parent Mode
	 */
	protected void notifyListeners() {
		// TODO could these be done earlier (or just once?)
		JMeterContext threadContext = getThreadContext();
		JMeterVariables threadVars = threadContext.getVariables();
		SamplePackage pack = (SamplePackage) threadVars.getObject(JMeterThread.PACKAGE_OBJECT);
		if (pack == null) {
			// If child of TransactionController is a ThroughputController and TPC does
			// not sample its children, then we will have this
			// TODO Should this be at warn level ?
			log.warn("Could not fetch SamplePackage");
		} else {
			SampleEvent event = new SampleEvent(res, threadContext.getThreadGroup().getName(), threadVars, true);
			// We must set res to null now, before sending the event for the transaction,
			// so that we can ignore that event in our sampleOccured method
			res = null;
			lnf.notifyListeners(event, pack.getSampleListeners());
		}
	}

	@Override
	public void sampleOccurred(SampleEvent se) {

		// TODO 此处先将SampleResult保存下来，待异步执行完毕之后再进行处理

		if (!isGenerateParentSample()) {
			// Check if we are still sampling our children
			if (res != null && !se.isTransactionSampleEvent()) {

				sampleResultList.add(se.getResult());
			}
		}
	}

	@Override
	public void sampleStarted(SampleEvent e) {

	}

	@Override
	public void sampleStopped(SampleEvent e) {

		log.info("sampleStopped...");
	}

	/**
	 * Whether to include timer and pre/post processor time in overall sample.
	 *
	 * @return boolean (defaults to true for backwards compatibility)
	 */
	public boolean isIncludeTimers() {
		return getPropertyAsBoolean(INCLUDE_TIMERS, DEFAULT_VALUE_FOR_INCLUDE_TIMERS);
	}

	/**
	 * Whether to include timers and pre/post processor time in overall sample.
	 *
	 * @param includeTimers Flag whether timers and pre/post processor should be included in overall sample
	 */
	public void setIncludeTimers(boolean includeTimers) {
		setProperty(INCLUDE_TIMERS, includeTimers, DEFAULT_VALUE_FOR_INCLUDE_TIMERS);
	}
}
